import torch
import torchvision
from torch import nn


class Generator_X(nn.Module): # the generator of set x
    
    def __init__(self, *args, **kwargs) -> None:
        super(Generator_X, self).__init__(*args, **kwargs)

    def forward(self, x):
        pass


class Generator_Y(nn.Module): # the generator of set y

    def __init__(self, *args, **kwargs) -> None:
        super(Generator_Y, self).__init__(*args, **kwargs)

    def forward(self, x):
        pass


# generator apartment
class Encoder(nn.Module): # encode

    def __init__(self, *args, **kwargs) -> None:
        super(Encoder, self).__init__(*args, **kwargs)
        # self.conv_1 = nn.Conv2d(3, 64, 7, 1, 3, padding_mode='reflect')
        # self.instanceNorm2d_1 = nn.InstanceNorm2d(64)
        # self.relu = nn.ReLU()
        # self.conv_2 = nn.Conv2d(64, 128, 3, 2, 1)
        # self.instanceNorm2d_2 = nn.InstanceNorm2d(128)
        # self.conv_3 = nn.Conv2d(128, 256, 3, 2, 1)
        # self.instanceNorm2d_3 = nn.InstanceNorm2d(256)
        self.main = nn.Sequential(*[
            nn.Conv2d(3, 64, 7, 1, 3, padding_mode='reflect'),
            nn.InstanceNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, 2, 1),
            nn.InstanceNorm2d(128),
            nn.Conv2d(128, 256, 3, 2, 1),
            nn.InstanceNorm2d(256),
        ])

    def forward(self, x):
        return self.main(x)


class ResidualBlock(nn.Module):
    def __init__(self, *args, **kwargs) -> None:
        super(ResidualBlock, self).__init__(*args, **kwargs)

        block = [
            nn.Conv2d(256, 256, 3, 1, 1, padding_mode='reflect'),
            nn.InstanceNorm2d(256),
            nn.ReLU(),
            nn.Conv2d(256, 256, 3, 1, 1, padding_mode='reflect'),
            nn.InstanceNorm2d(256),
        ]

        self.block = nn.Sequential(*block)

    def forwardd(self, x):
        return x + self.block(x)
    

class Decoder(nn.Module): # decode
    
    def __init__(self, *args, **kwargs) -> None:
        super(Decoder, self).__init__(*args, **kwargs)
        self.main = nn.Sequential(*[
            nn.ConvTranspose2d(256, 128, 3, 2, 1, output_padding=1),
            nn.InstanceNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 3, 2, 1, output_padding=1),
            nn.InstanceNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 3, 7, 1, 3, padding_mode='reflect'),
            nn.Tanh(),
        ])

    def forward(self, x):
        return self.main(x)
